import os
from PIL import Image
from torchvision import transforms
import random
import pandas as pd

def augment_dataframe_to_target(df, label_col, target_count, img_dir, save_dir=None, transform=None):
    """
    对DataFrame进行数据增强，生成增强图片并更新DataFrame

    Args:
        df: 包含图片文件名和标签的DataFrame
        label_col: 标签列名（如'diagnosis'）
        target_count: 目标总数
        img_dir: 原始图片目录
        save_dir: 增强图片保存目录（默认与img_dir相同）
        transform: 自定义变换，默认为随机翻转+旋转+裁剪

    Returns:
        增强后的DataFrame（包含原始+增强数据）
    """
    if transform is None:
        transform = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomVerticalFlip(),
            transforms.RandomRotation(30),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1),
            transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0))
        ])


    if save_dir is None:
        save_dir = img_dir

    os.makedirs(save_dir, exist_ok=True)

    current_count = len(df)
    if current_count >= target_count:
        return df

    augmented_rows = []
    for i in range(target_count - current_count):
        # 随机选择一行
        row = df.sample(n=1).iloc[0]
        img_path = os.path.join(img_dir, row['id_code'])
        img = Image.open(img_path)

        # 增强并保存
        aug_img = transform(img)
        new_name = f"aug_{i}_{row['id_code']}"
        aug_img.save(os.path.join(save_dir, new_name))

        # 添加新行
        augmented_rows.append({'id_code': new_name, label_col: row[label_col]})

    return pd.concat([df, pd.DataFrame(augmented_rows)], ignore_index=True)